#include <cmath>
#include <iostream>
#include <cstdio>

#include "Triangle.h"
#include "Types.h"
#include "FEMModel.h"
#include "Node.h"
#include "Tetrahedron.h"
#include "DeformationState.h"
#include "BigVector.h"
#include "Setting.h"

using namespace std;

void UniformDilatationSolution (FiniteElementModel*fem,int axis, Real);
void UniformDilatationFix (FiniteElementModel*fem, int axis,Real f);
void UniformDilatationParameters (Real *elongate, Real *contract, Real f);
void UniformCompressionFix (FiniteElementModel*fem, Real f);
void UniformCompressionDisplacement (FiniteElementModel*fem, Real f);

void FixNodes (FiniteElementModel*fem, Array<int> fs)
{
	for (int i=fs.Size ();i--;)
		fem->FixNode ( fem->GetNode (fs[i]), true);
}

Array<int> FindOutsideFace (FiniteElementModel* fem, int axis, int sign)
{
	Array<int> idxes;
	Vector3 bounds[3];

	NodesBoundingBox (fem, false, &bounds[0],&bounds[2]);
	for (int i= fem->NodeCount (); i--;)
	{
		Vector3 v = fem->GetNode (i)->Location (false);

		if (fabs (v(axis) - bounds[sign + 1](axis)) < 1e-4)
			idxes.Push (i);
	}
	return idxes;
}

void FixOutside(FiniteElementModel* fem, int axis, int sign)
{
	Array<int> idxes = FindOutsideFace (fem, axis, sign);
	FixNodes (fem,idxes);  
}

/*
fix backside, dilate horizontally. 
*/
void InitialForceOnCubeFixedBack (FiniteElementModel*fem, Real force)
{
	Array<int> left;
	Array<int> right;  
	Array<int> back;

	fprintf (stderr, "Fixing back, pulling left and right.");
	for (int i=left.Size ();i--;)
		for (int j =0; j  < DIM ; j++)
		{
			fem->AddForce (fem->GetNode (i), Vector3 (-force,0,0));
		}
		for (int i=right.Size ();i--;)
			for (int j =0; j  < DIM ; j++)
			{
				fem->AddForce (fem->GetNode (i), Vector3 (force,0,0));
			}

			FixNodes (fem,back);
}

void RotatedForceOnFixedBack (FiniteElementModel*fem, Real force)
{
	Real angle = GetNumberSetting ("init-force-angle");

	Array<int> left;
	Array<int> right;  
	Array<int> back;
	// find_lrb (&left,&right,&back,fem);

	cerr << " fixing back, applying force of "<<force<<" at " << angle << " pi\n";
	angle *= M_PI;

	LinkArray<Node> outer_nodes ;

	set<Triangle*>::const_iterator bt (fem->OuterFacesBegin ());
	set<Triangle*>::const_iterator et (fem->OuterFacesEnd ());
	for (; bt != et ; bt++)
		for (int j=3 ; j--;)
			outer_nodes.Push ((*bt)->GetNode (j));

	outer_nodes.DefaultSort ();
	outer_nodes.Uniq();
	Vector3 force_vec (cos (angle), sin (angle), 0.0);

	Real total_force = 0.;
	Array<Vector3> dfx;
	Array<int> idx;
	for (int i= 0; i < outer_nodes.Size (); i++)
	{
		Vector3 v = outer_nodes[i]->Location (false);
		idx.Push (outer_nodes[i]->Number () );

		Real radangle =  acos (force_vec * v / (v.Length ()));

		Real f1 = exp ( -3 * Sqr (radangle));
		Real f2 = exp ( -3 * Sqr (M_PI- radangle));

		Vector3 df =f1 * force_vec  - f2 *force_vec;
		dfx.Push (df);
		total_force   += df.Length ();   
	}

	for (int i =0; i < dfx.Size (); i++)
		for (int j =0; j < 3; j++)
			fem->AddForce (fem->GetNode (i), dfx[i] * (force / total_force));

	FixNodes (fem,back);
}


void InitialForceOnCubeSymGaussian (FiniteElementModel*fem, int axis, Real force)
{
	cerr << "Symmetrical gaussian force."<<endl;   
	Array<int> left = FindOutsideFace (fem, axis, -1);
	Array<int> right = FindOutsideFace (fem, axis, 1);  
	Array<int> *ps[] = {&left, &right};

	int d = 1;
	Array<Vector3> vs;
	vs.SetSize (fem->NodeCount ());

	Real total[] = {0,0};
	for (int k = 0 ; k < 2; k++)
	{    int sz =ps[k]->Size ();

	for (int i= sz; i--; )
	{
		int j = ps[k]->Elem (i);

		Vector3 l=  fem->GetNode (j)->Location(false);
		l(axis) = 0.0;

		Real fact = 2 * sqrt(M_PI)*16*exp (-Sqr (4*l.Length ()));

		Real axforce = fact;
		total[k] += axforce;
		Vector3 f;
		f(axis) = (2 *k -1)* axforce;
		vs[j] = f;
	}
	}

	for (int k = 0; k< 2; k++)
	{
		int sz =ps[k]->Size ();
		Real tot = total[k];	
		for (int i= sz; i--; )
		{
			int j = ps[k]->Elem (i);
			Vector3 vf = vs[j];
			vf *= force / tot;
			fem->AddForce (fem->GetNode (j), vf);
		}
	} while ((d *= -1) != 1);
}


void InitialForceOnCubeSymUniform (FiniteElementModel*fem, int axis, Real ext_force)
{
	cerr << "Symmetrical uniform force."<<endl;   

	Array<int> left = FindOutsideFace (fem, axis, -1);
	Array<int> right = FindOutsideFace (fem, axis, 1);  
	Array<int> *ps[] = {&left, &right};

	Vector3 bounds[3];
	NodesBoundingBox (fem, false, &bounds[0],&bounds[2]);

	cout  << bounds[0] << " " << bounds[2];

	Array<Vector3> fvecs;
	fvecs.SetSize (fem->NodeCount ());
	Real ftotals[3] ={ 0,0,0};
	int d = 1;
	do {
		int sz =ps[(d+1)/2]->Size ();
		for (int i= sz; i--; )
		{
			int j = ps[(d+1)/2]->Elem (i);


			int boundary = 0;
			Vector3 l=  fem->GetNode (j)->Location(false);
			for (int k = 0; k< 3; k++) 
				if (axis != k
					&& (fabs (l(k) - bounds[0](k)) < 1e-4
					||fabs (l(k) - bounds[2](k)) < 1e-4))
					boundary ++;

			Vector3 fvec;
			Real f = (4 >> boundary ) / 4.0;
			ftotals[(d+1)/2] += f;
			fvec(axis) = d*f;

			fvecs[j] = fvec;
		}
	} while ((d *= -1) != 1);

	do {
		int sz =ps[(d+1)/2]->Size ();
		for (int i= sz; i--; )
		{
			int j = ps[(d+1)/2]->Elem (i);

			Vector3 fvec =fvecs[j];
			fvec *= ext_force /  ftotals[(d+1)/2];
			fem->AddForce (fem->GetNode (j), fvec);
		}
	} while ((d *= -1) != 1);
}


void ForceOnSingleFace (FiniteElementModel*fem, int axis, int sign,Real force)
{
	Array<int> face =  FindOutsideFace (fem,axis,sign);

	fem->FixNode (fem->GetNode (0),true);
	for (int i=face.Size ();i--;)
	{
		fem->AddForce (fem->GetNode (face[i]), Vector3 (force/face.Size (),0,0));
	}
}

void InitialForceAsymGaussian (FiniteElementModel*fem, int axis, int sign,  Real force)
{
	Array<int> is =   FindOutsideFace (fem,axis,sign);

	for (int i=is.Size ();i--;)
	{
		Vector3 l=  fem->GetNode (is[i])->Location(false);
		l(axis) = 0.0;

		Real fact = 2 * sqrt(M_PI)*16*exp (-Sqr (4*l.Length ()));
		fem->AddForce (fem->GetNode (is[i]), Vector3 (force / is.Size () *  fact, 0,0));
	}
}

void FixSymmetrical (FiniteElementModel*fem)
{
	printf ("Symmetric fixing\n");
	for (int i= fem->NodeCount (); i--;)
	{
		Vector3 v = fem->GetNode (i)->Location (false);
		for (int j = 0; j < 3; j++)
			if (fabs (v(j)) < 1e-4)
			{
				fem->pstate_->fixed_arr_[i*3 + j] = true;
			}
	}
}




void ForceOnCenterPoint (FiniteElementModel*fem, Axes a, int dir, Real force)
{
	printf ("Stretching force pulse of %lf\n", force);

	Array<int> face = FindOutsideFace (fem, a, dir);

	Vector3 fv;
	fv (a) = dir * force;

	for (int j = face.Size (); j--; )
	{
		Vector3 l = fem->GetNode (face[j])->Location (false);
		l(a) = 0.0;
		if (fabs (l.Length ()) < 1e-8) 
			fem->AddForce (fem->GetNode (face[j]), fv);
	}

}



void BendForceOnCenterPoint (FiniteElementModel*fem, Axes a, int dir, Real force)
{
	printf ("Bending force pulse of %lf\n", force);
	Array<int> face = FindOutsideFace (fem, a, dir);

	Vector3 fv;
	fv ((a+1)%3) = dir * force;

	for (int j = face.Size (); j--; )
	{
		Vector3 l = fem->GetNode (face[j])->Location (false);
		l(a) = 0.0;
		if (fabs (l.Length ()) < 1e-8)
		{
			fem->AddForce (fem->GetNode (face[j]), fv);
		}
	}
}

void GravityForce (FiniteElementModel*fem, Real acc)
{
	Real d = GetNumberSetting ("density") * acc ;

	Array<Real> lv =LumpedVolumes (fem);

	Real f = 0.0;
	for (int i = lv.Size(); i--;)
	{
		Real df = lv[i]*d;
		fem->AddForce (fem->GetNode (i),Vector3 (0,-df, 0));
		f += df;
	}

	printf ("Total gravity force:  %lf\n" , f);
}

void SetInitialForce (FiniteElementModel*fem)
{
	if (Real g = GetNumberSetting ("gravity"))
	{
		GravityForce (fem, g);
	}

	const char * fix_shape = GetStringSetting ("init-fix-shape");
	Real f = GetNumberSetting ("init-force");

	if (0)
		;
	else if (!strcmp (fix_shape, "bottom"))
		FixOutside (fem, Y_AXIS, -1);
	else if (!strcmp (fix_shape, "back"))
		FixOutside (fem, Z_AXIS, -1);
	else if (!strcmp (fix_shape, "left"))
		FixOutside (fem, X_AXIS, -1);
	else if (!strcmp (fix_shape, "right"))
		FixOutside (fem, X_AXIS, 1);
	else if (!strcmp (fix_shape, "symmetrical"))
		FixSymmetrical (fem);
	else if (!strcmp (fix_shape, "dilation"))
	{
		FixSymmetrical (fem);
		UniformDilatationFix (fem, X_AXIS, f);
		return;
	}
	else if (!strcmp (fix_shape, "expansion"))
	{
		FixSymmetrical (fem);
		UniformCompressionDisplacement (fem, f);
		return;
	}

	if(!f)
		return ;


	const char * shape = GetStringSetting ("init-force-shape");

	if (!strcmp (shape, "none"))
		;
	else if (!strcmp (shape, "pulse"))
		ForceOnCenterPoint (fem, Z_AXIS, 1, f);
	else if (!strcmp (shape, "bend-pulse"))
		BendForceOnCenterPoint (fem, Z_AXIS, 1, f);
	else if (!strcmp (shape, "right"))
		ForceOnSingleFace (fem, X_AXIS, 1, f);
	else if (!strcmp (shape, "sym-gauss"))
		InitialForceOnCubeSymGaussian (fem, X_AXIS, f);
	else if (!strcmp (shape, "sym-uniform"))
		InitialForceOnCubeSymUniform (fem, X_AXIS, f);

	cout << "||Init force||_1 " << BigVector1Norm (fem->pstate_->delta_external_force_.AccessArray (),
		fem->pstate_->Dimension())<<endl;


	if (!strcmp (fix_shape, "dilation")
		&& !strcmp (GetStringSetting ("init-solution"), "yes"))
		UniformDilatationSolution (fem,X_AXIS, f);
}


/*
The solution for uniform dilatation along AXIS. 
*/
void UniformDilatationSolution (FiniteElementModel*fem, int axis, Real f)
{
	Real contract, elongate;
	cout << "setting uniform dilation solution " << endl;

	UniformDilatationParameters (&elongate, &contract, f);
	Real const*  locs = fem->node_locations_.AccessArray ();
	for (int n = fem->NodeCount (); n--; )
		for (int a = 0; a < 3; a++)
		{
			fem->pstate_->displacement_[3*n + a] =
				((a == axis) ? elongate : -contract) * locs [3*n + a];
		}
}

/*
Set boundary conditions for  uniform (ie. fix the outside) compression. 
*/
void UniformCompression (FiniteElementModel*fem, Real f)
{
	Real const *locs = fem->node_locations_.AccessArray ();
	Real *disp = fem->pstate_->displacement_.UnsafeAccessArray();
	bool *fix= fem->pstate_->fixed_arr_.UnsafeAccessArray ();

	Real factor = f;
	for (int axis =0; axis < 3; axis++)
	{
		int d = -1;
		do {
			Array<int> face = FindOutsideFace (fem, axis, d);
			for (int j = face.Size (); j--; )
			{
				int k = face[j]*3 + axis;

				disp[k] = factor * locs[k];
				if (axis == X_AXIS)
					disp[k] *= -1;
				fix[k] = 1;
			}

		}while ((d*=-1)!=-1);
	}
}

/*
Compress the model uniformly, throughout the volume.
*/
void UniformCompressionDisplacement (FiniteElementModel*fem, Real f)
{
	Real const *locs = fem->node_locations_.AccessArray ();
	Real *disp = fem->pstate_->displacement_.UnsafeAccessArray();


	Real factor = f;

	for (int n = fem->NodeCount (); n--;)
	{
		for (int axis =0; axis < 3; axis++)
		{
			int k = n*3 + axis;

			disp[k] = factor * locs[k];
			if (axis == X_AXIS)
				disp[k] *= -2;
		}
	}
}

/*
Fix the outside faces of AXIS to correspond with uniform dilation in
AXIS direction.  */
void UniformDilatationFix (FiniteElementModel*fem, int axis, Real f)
{
	Real contract, elongate;
	cout << "Fixing sides for uniform dilation " << endl;

	UniformDilatationParameters (&elongate, &contract, f);

	Array<int> left = FindOutsideFace (fem, X_AXIS, -1);
	Array<int> right = FindOutsideFace (fem, X_AXIS, 1); 

	Real const *locs = fem->node_locations_.AccessArray ();
	Real *disp = fem->pstate_->displacement_.UnsafeAccessArray();

	left.Concat (right);
	for (int n = left.Size(); n--;)
	{
		int j = 3*left[n] + axis;

		disp[j] = elongate * locs [j];
		fem->pstate_->fixed_arr_[j] = 1;
	}

}

void UniformDilatationParameters (Real *elongate, Real *contract, Real f)
{
	Real young = GetNumberSetting ("young");
	Real poisson = GetNumberSetting ("poisson");

	cout << "Young, Poisson = " << young << " "  << poisson << "\n";

	bool linear = !strcmp (GetStringSetting ("elasticity"), "linear");

	cout <<  ((linear)  ? "linear " : "nonlinear ") << " elasticity"<<endl;


	/*
	TODO: check paper!
	*/
	if (linear)
		*elongate = f / young;
	else
		*elongate = (-1 + sqrt (1 + 2 * f  / young));

	Real discr = 1 - (2 * f * poisson / young);
	if (discr < 0)
	{
		cout << "No exact solution??"<<endl;
		return ;
	}

	if (linear)
		*contract = (f * poisson / young);
	else
		*contract= (1 - sqrt (discr));

	cout << "Uniform dilatation.  Contraction: " << *contract << " elongation: " << *elongate << endl;

}